#include <Windows.h>
#include <winnt.h>
#include <stdio.h>
#include "MemoryModule.h"

#pragma warning( disable : 4311 4312 )


typedef BOOL (WINAPI *DllEntryProc)(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved);

#define GET_HEADER_DICTIONARY(headers, idx) &(headers)->OptionalHeader.DataDirectory[idx]


#ifndef IMAGE_SIZEOF_BASE_RELOCATION
// Vista SDKs no longer define IMAGE_SIZEOF_BASE_RELOCATION!?
#define IMAGE_SIZEOF_BASE_RELOCATION (sizeof(IMAGE_BASE_RELOCATION))
#endif

static void CopySections(void * code, const void * data, PIMAGE_NT_HEADERS nt_header)
{
    int i;
    unsigned char *dest;
    unsigned char *codeBase = (unsigned char *)code;
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(nt_header);

    for(i = 0; i < nt_header->FileHeader.NumberOfSections; i++, section++)
    {
        dest = codeBase + section->VirtualAddress;
        memcpy(dest, (unsigned char *)data + section->PointerToRawData, section->SizeOfRawData);
        section->Misc.PhysicalAddress = (DWORD)dest;
    }
}



static void PerformBaseRelocation(void * code, PIMAGE_NT_HEADERS nt_header, DWORD delta)
{
    DWORD i;
    unsigned char *codeBase = (unsigned char *)code;

    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(nt_header, IMAGE_DIRECTORY_ENTRY_BASERELOC);

    if(directory->Size > 0)
    {
        PIMAGE_BASE_RELOCATION relocation = (PIMAGE_BASE_RELOCATION)(codeBase + directory->VirtualAddress);

        for(; relocation->VirtualAddress > 0;)
        {
            unsigned char *dest = (unsigned char *)(codeBase + relocation->VirtualAddress);
            unsigned short *relInfo = (unsigned short *)((unsigned char *)relocation + IMAGE_SIZEOF_BASE_RELOCATION);

            for(i = 0; i < ((relocation->SizeOfBlock - IMAGE_SIZEOF_BASE_RELOCATION) / 2); i++, relInfo++)
            {
                DWORD *patchAddrHL;
                int type, offset;

                // the upper 4 bits define the type of relocation
                type = *relInfo >> 12;
                // the lower 12 bits define the offset
                offset = *relInfo & 0xfff;

                switch(type)
                {
                    case IMAGE_REL_BASED_ABSOLUTE:
                        // skip relocation
                        break;

                    case IMAGE_REL_BASED_HIGHLOW:
                        // change complete 32 bit address
                        patchAddrHL = (DWORD *)(dest + offset);
                        *patchAddrHL += delta;
                        break;

                    default:
                        //printf("Unknown relocation: %d\n", type);
                        break;
                }
            }

            // advance to next relocation block
            relocation = (PIMAGE_BASE_RELOCATION)(((DWORD)relocation) + relocation->SizeOfBlock);
        }
    }
}

static int BuildImportTable(void * code, PIMAGE_NT_HEADERS nt_header)
{
    int result = 1;
    unsigned char *codeBase = (unsigned char *)code;

    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(nt_header, IMAGE_DIRECTORY_ENTRY_IMPORT);

    if(directory->Size > 0)
    {
        PIMAGE_IMPORT_DESCRIPTOR importDesc = (PIMAGE_IMPORT_DESCRIPTOR)(codeBase + directory->VirtualAddress);

        for(; importDesc->Name; importDesc++)
        {
            DWORD *thunkRef, *funcRef;
            HMODULE handle = GetModuleHandleA((LPCSTR)(codeBase + importDesc->Name));
			
			if (handle == NULL)
			{
				handle = LoadLibraryA((LPCSTR)(codeBase + importDesc->Name));
			}

            if(handle == INVALID_HANDLE_VALUE)
            {
                result = 0;
                break;
            }


            if(importDesc->OriginalFirstThunk)
            {
                thunkRef = (DWORD *)(codeBase + importDesc->OriginalFirstThunk);
                funcRef = (DWORD *)(codeBase + importDesc->FirstThunk);
            }
            else
            {
                // no hint table
                thunkRef = (DWORD *)(codeBase + importDesc->FirstThunk);
                funcRef = (DWORD *)(codeBase + importDesc->FirstThunk);
            }

            for(; *thunkRef; thunkRef++, funcRef++)
            {
                if IMAGE_SNAP_BY_ORDINAL(*thunkRef)
                {
                    *funcRef = (DWORD)GetProcAddress(handle, (LPCSTR)IMAGE_ORDINAL(*thunkRef));
                }
                else
                {
                    PIMAGE_IMPORT_BY_NAME thunkData = (PIMAGE_IMPORT_BY_NAME)(codeBase + *thunkRef);
                    *funcRef = (DWORD)GetProcAddress(handle, (LPCSTR)&thunkData->Name);
                }

                if(*funcRef == 0)
                {
                    result = 0;
                    break;
                }
            }

            if(!result)
            {
                break;
            }
        }
    }

    return result;
}

// Protection flags for memory pages (Executable, Readable, Writeable)
static int ProtectionFlags[2][2][2] =
{
    {
        // not executable
        {PAGE_NOACCESS, PAGE_WRITECOPY},
        {PAGE_READONLY, PAGE_READWRITE},
    }, {
        // executable
        {PAGE_EXECUTE, PAGE_EXECUTE_WRITECOPY},
        {PAGE_EXECUTE_READ, PAGE_EXECUTE_READWRITE},
    },
};

static void FinalizeSections(void * code, PIMAGE_NT_HEADERS nt_header)
{
    int i;
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(nt_header);

    // loop through all sections and change access flags
    for(i = 0; i < nt_header->FileHeader.NumberOfSections; i++, section++)
    {
        DWORD protect, oldProtect, size;
        int executable = (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;
        int readable = (section->Characteristics & IMAGE_SCN_MEM_READ) != 0;
        int writeable = (section->Characteristics & IMAGE_SCN_MEM_WRITE) != 0;

        // determine protection flags based on characteristics
        protect = ProtectionFlags[executable][readable][writeable];

        if(section->Characteristics & IMAGE_SCN_MEM_NOT_CACHED)
        {
            protect |= PAGE_NOCACHE;
        }

        // determine size of region
        size = section->SizeOfRawData;

        if(size == 0)
        {
            if(section->Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA)
            {
                size = nt_header->OptionalHeader.SizeOfInitializedData;
            }
            else if(section->Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)
            {
                size = nt_header->OptionalHeader.SizeOfUninitializedData;
            }
        }

        if(size > 0)
        {
            // change memory access flags
            if(VirtualProtect((LPVOID)section->Misc.PhysicalAddress, section->SizeOfRawData, protect, &oldProtect) == 0)
            {
            }
        }
    }

}


static int GetExportFunction(void * code, PIMAGE_NT_HEADERS nt_header, void **out_addr)
{
    int result = 0;
    unsigned char *codeBase = (unsigned char *)code;

    if(out_addr == NULL)
    {
        return 1;
    }

    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(nt_header, IMAGE_DIRECTORY_ENTRY_EXPORT);

    if(directory->Size == 0)
    {
        return result;
    }

    PIMAGE_EXPORT_DIRECTORY exportDesc = (PIMAGE_EXPORT_DIRECTORY)(codeBase + directory->VirtualAddress);

    if(exportDesc->NumberOfFunctions != 1)
    {
        return result;
    }

    *(DWORD *)out_addr = DWORD(codeBase + * (DWORD *)(codeBase + exportDesc->AddressOfFunctions));

    result = 1;
    return result;
}

void* MemoryLoadLibrary(const void *data, void **out_addr)
{
    PIMAGE_DOS_HEADER dos_header;
    PIMAGE_NT_HEADERS nt_header;
    unsigned char *code;
    DWORD locationDelta;
    DllEntryProc DllEntry;
    BOOL successfull;
    int len = 0;
    DWORD entry_point = 0;


    dos_header = (PIMAGE_DOS_HEADER)data;
    nt_header = (PIMAGE_NT_HEADERS) & ((const unsigned char *)(data))[dos_header->e_lfanew];

    code = (unsigned char *)VirtualAlloc(NULL,
                                         nt_header->OptionalHeader.SizeOfImage,
                                         MEM_RESERVE | MEM_COMMIT,
                                         PAGE_EXECUTE_READWRITE);

    if(code == NULL)
    {
        return NULL;
    }

    memcpy(code, dos_header, dos_header->e_lfanew + nt_header->OptionalHeader.SizeOfHeaders);


    nt_header = (PIMAGE_NT_HEADERS) & ((const unsigned char *)(code))[dos_header->e_lfanew];

 
    CopySections(code, data, nt_header);


    locationDelta = (DWORD)(code - nt_header->OptionalHeader.ImageBase);

    if(locationDelta != 0)
    {
        PerformBaseRelocation(code, nt_header, locationDelta);
    }

    if(!BuildImportTable(code, nt_header))
    {
        goto error;
    }

    if(!GetExportFunction(code, nt_header, out_addr))
    {
        goto error;
    }


    len = nt_header->OptionalHeader.SizeOfImage;
    entry_point = nt_header->OptionalHeader.AddressOfEntryPoint;

    FinalizeSections(code, nt_header);


    if(entry_point != 0)
    {
        DllEntry = (DllEntryProc)(code + entry_point);

        if(DllEntry == 0)
        {
            goto error;
        }

        successfull = (*DllEntry)((HINSTANCE)code, DLL_PROCESS_ATTACH, (LPVOID)NULL);

        if(!successfull)
        {
            goto error;
        }

    }

    return code;
error:
    VirtualFree(code, 0, MEM_RELEASE);
    return NULL;
}

void MemoryFreeLibrary(void *data)
{
    VirtualFree(data, 0, MEM_RELEASE);
}
